import pylab
import nibabel
import numpy
import scipy
import math
import matplotlib as plt

# Matrices

# Traslacion
def T(dx, dy, dz):
    res = numpy.identity(4, dtype = float)
    res[0][3] = dx
    res[1][3] = dy
    res[2][3] = dz    
    return res

# Rotacion
def R(tx, ty, tz):
    r_x = numpy.identity(4, dtype = float)
    r_x[1][1] = math.cos(tx)
    r_x[2][2] = r_x[1][1]
    r_x[2][1] = math.sine(tx)
    r_x[1][2] = -1 * r_x[2][1]
    
    r_y = numpy.identity(4, dtype = float)
    r_y[0][0] = math.cos(ty)
    r_y[2][2] = r_y[0][0]
    r_y[2][0] = math.sine(ty)
    r_y[0][2] = -1 * r_y[2][0]
    
    r_z = numpy.identity(4, dtype = float)
    r_z[0][0] = math.cos(tz)
    r_z[1][1] = r_z[0][0]
    r_z[1][0] = math.sine(tz)
    r_z[0][1] = -1 * r_z[1][0]
        
    return r_x * r_y * r_z


# Escalamiento (usual)
def S_usual(sx, sy, sz):
    res = numpy.identity(4, dtype = float)
    res[0][0] = sx
    res[1][1] = sy
    res[2][2] = sz    
    return res


# Escalamiento anisotropico (shear)
def S(sxy, sxz, syx, syz, szx, szy):
    res = numpy.identity(4, dtype = float)
    res[0][1] = sxy
    res[0][2] = sxz
    res[1][0] = syx
    res[1][2] = syz
    res[2][0] = szx
    res[2][2] = szy
    return res


# Convolucion

def convolve(img, kernel, n):    
    mid = n / 2
    
    l = numpy.zeros((img.shape[0], img.shape[1], img.shape[2]), dtype = float)
    
    for i in xrange(mid, img.shape[0] - mid):
        for j in xrange(mid, img.shape[1] - mid):
            for k in xrange(mid, img.shape[2] - mid):                
                l[i, j, k] = conv(img, kernel, i, j, k)
    return l


def getSubmatrix(matrix, init, size):
    x = init[0]
    y = init[1]
    z = init[2]
    
    s_x = size[0]
    s_y = size[1]
    s_z = size[2]
    
    return matrix[x - s_x : x + s_x + 1 , y - s_y : y + s_y + 1, z - s_z : z + s_z + 1]

def conv(img, window, i, j, k):    
    window_size = window.shape
    img_size = img.shape
    w_x = window_size[0] / 2
    w_y = window_size[1] / 2
    w_z = window_size[2] / 2
	
    #Es un borde, entonces queda igual
    if (i < w_x or j < w_y or k < w_z):
        return img[i, j, k]
	
    if((img_size[0] - 1) - i < w_x or (img_size[1] - 1) - j < w_y or (img_size[2] - 1) - k < w_z):
        return img[i, j, k]
	
	#print(i, j, k)
    subimg = getSubmatrix(img, [i, j, k], [w_x, w_y, w_z])
	
    return numpy.multiply(subimg, window).sum() / window.sum()

    
# Transformaciones
# Parametros:
# Imagen original + parametros

def transform_traslation(I, dx, dy, dz):
    trans = T(dx, dy, dz)
    print trans, I
    return convolve(trans, I, 4)

def transform_rotation(I, tx, ty, tz):
    trans = R(tx, ty, tz)
    return trans * I
    
def transform_rigid(I, dx, dy, dz, tx, ty, tz):
    trans = R(tx, ty, tz) * T(dx, dy, dz)
    return trans * I

def transform_scaling(I, sxy, sxz, syx, syz, szx, szy):
    trans = S(sxy, sxz, syx, syz, szx, szy)
    return trans * I
    
def transform_affine(I, dx, dy, dz, tx, ty, tz, sxy, sxz, syx, syz, szx, szy):
    trans = S(sxy, sxz, syx, syz, szx, szy) * R(tx, ty, tz) * T(dx, dy, dz)
    return trans * I
    

# Registraciones
# Parametros:
# Imagen original + fuente
# Se minimiza la energia usando cuadrados minimos

def registration_traslation(I, F):
    v, success = scipy.optimize.leastsq(e(I, transform_traslation(F, dx, dy, dz)), 0, args = (dx, dy, dz)) 
    return v
    
def e(l, i):
    return sim(i, l) + reg(l)
    
def sim(i, l):
    v_i = i.flatten()
    v_l = l.flatten()
    
    r = v_i - v_l

    return multiply(r, r).sum() / len(r)
    
def reg(j):
    v_j = j.flatten()
    return absolute(gradient(v_j)).sum() / len(v_j)
